using BioSequences
using StaticArrays
using BenchmarkTools

const cmatrix = [
    0 1 2 1 3
    1 0 1 5 1
    2 1 0 9 1
    1 5 9 0 1
    3 1 1 1 0
];
const basesnum = UInt8.(['B', 'A', 'T', 'G', 'C']) .- UInt8('@')
const tmpmat = zeros(Int, 20, 20)
tmpmat[basesnum, basesnum] = cmatrix
const dist = SMatrix{20,20}(tmpmat)

const bias = UInt8('A') - UInt8(1)
const a = readline("data/Data-P2a.txt");
const b = readline("data/Data-P2b.txt");
const Seq_A = Base.CodeUnits(a) .- bias
const Seq_B = Base.CodeUnits(b) .- bias

function dp!(s1, e1, s2, e2)
    if s1 == e1
        rsA, rsB = LongDNASeq("-"^(e2 - s2)), LongDNASeq(b[s2+1:e2])
        cost = sum(dist[2, Seq_B[s2+1:e2]])
        return rsA, rsB, cost
    elseif  s2 == e2
        rsA, rsB = LongDNASeq(a[s1+1:e1]), LongDNASeq("-"^(e1 - s1))
        cost = sum(dist[Seq_A[s1+1:e1], 2])
        return rsA, rsB, cost
    elseif s1 == e1 - 1
        strt1, strt2 = dist[2, Seq_B[s2+1:e2]], dist[Seq_A[e1], 2]
        save = strt1 .+ strt2 - dist[Seq_A[e1], Seq_B[s2+1:e2]]
        rsA, rsB = LongDNASeq("-"^(e2 - s2)), LongDNASeq(b[s2+1:e2])
        pos = argmax(save)
        cost = sum(strt1) + strt2
        if save[pos] > 0 
            rsA[pos] = DNA(a[e1])
            cost -= save[pos]
        else
            push!(rsA, DNA(a[e1]))
            push!(rsB, DNA_Gap)
        end
        return rsA, rsB, cost
    end
    f1 = zeros(Int, e2 - s2 + 1)
    f2, g1 = f1[:], f1[:]
    
    mid = (e1 + s1) ÷ 2
    f1[1] = 0
    f1[2:e2-s2+1] = cumsum(dist[2, Seq_B[s2+1:e2]])
    @inbounds for i = s1+1:mid
        f2[1] = f1[1] + dist[Seq_A[i], 2]
        for j = s2+2:e2+1
            f2[j-s2] = min(
                f1[j-s2] + dist[Seq_A[i], 2],
                f2[j-s2-1] + dist[2, Seq_B[j-1]],
                f1[j-s2-1] + dist[Seq_A[i], Seq_B[j-1]],
            )
        end
        f1, f2 = f2, f1
    end

    g2 = f2
    g1[e2-s2+1] = 0
    g1[e2-s2:-1:1] = cumsum(reverse(dist[2, Seq_B[s2+1:e2]]))
    @inbounds for i = e1:-1:mid+1
        g2[e2-s2+1] = g1[e2-s2+1] + dist[Seq_A[i], 2]
        for j = e2:-1:s2+1
            g2[j-s2] = min(
                g1[j-s2] + dist[Seq_A[i], 2],
                g2[j-s2+1] + dist[2, Seq_B[j]],
                g1[j-s2+1] + dist[Seq_A[i], Seq_B[j]],
            )
        end
        g1, g2 = g2, g1
    end
    pos = argmin(f1 + g1) + s2 - 1
    f1, f2, g1, g2 = nothing, nothing, nothing, nothing
    res1 = dp!(s1, mid, s2, pos)
    res2 = dp!(mid, e1, pos, e2)
    return res1[1] * res2[1], res1[2] * res2[2], res1[3] + res2[3]
end

function main()
    resA, resB, cost = dp!(0, length(a), 0, length(b))
    open("outA.txt", "w") do io
        write(io, String(resA))
    end
    open("outB.txt", "w") do io
        write(io, String(resB))
    end
    cost
end

@time main()
